﻿#pragma once
#include<iostream>
#include<assert.h>
#include<vector>
using namespace std;

template<class T>
struct AVLTreeNode
{
	AVLTreeNode(const T& data = T())
		: _pLeft(nullptr)
		, _pRight(nullptr)
		, _pParent(nullptr)
		, _data(data)
		, _bf(0)
	{}

	AVLTreeNode<T>* _pLeft;
	AVLTreeNode<T>* _pRight;
	AVLTreeNode<T>* _pParent;
	T _data;

	int _bf;   // 节点的平衡因子
};


// AVL: 二叉搜索树 + 平衡因子的限制
template<class T>
class AVLTree
{
	typedef AVLTreeNode<T> Node;
public:
	AVLTree()
		: _pRoot(nullptr)
	{}
	// 在AVL树中插入值为data的节点
	bool Insert(const T& data)
	{
		if (_pRoot == nullptr)
		{
			_pRoot = new Node(data);
			return true;
		}

		Node* parent = nullptr;
		Node* cur = _pRoot;

		// 按搜索树的规则插入节点
		while (cur)
		{
			if (cur->_data < data)
			{
				parent = cur;
				cur = cur->_pLeft;
			}
			else if (cur->_data > data)
			{
				parent = cur;
				cur = cur->_pRight;
			}
			else
			{
				return false;
			}
		}
		cur = new Node(data);
		if (parent->_data < data)
		{
			parent->_pRight = cur;
		}
		else
		{
			parent->_pLeft = cur;
		}
		cur->_pParent = parent;

		// 更新平衡因子
		while (parent)
		{
			// 更新平衡因子
			if (cur == parent->_pLeft)
				parent->_bf--;
			else
				parent->_bf++;

			// 调整
			if (parent->_bf == 0)
			{
				break;
			}
			else if (parent->_bf == -1 || parent->_bf == 1)
			{
				// 继续向上更新
				cur = parent;
				parent = cur->_pParent;
			}
			else if(parent->_bf == -2 || parent->_bf == 2)
			{
				// 旋转
				if (parent->_bf == -2 && cur->_bf == -1)  // 右单旋
				{
					RotateR(parent);
				}
				else if (parent->_bf == 2 && cur->_bf == 1)  // 左单旋
				{
					RotateL(parent);
				}
				else if (parent->_bf == -2 && cur->_bf == -1)  // 左右双旋
				{
					RotateLR(parent);
				}
				else if (parent->_bf == 2 && cur->_bf == -1)  // 右左双旋
				{
					RotateRL(parent);
				}
				else
				{
					assert(false);
				}
				break;
			}
			else
			{
				assert(false);
			}
		}
		return true;
	}

	// AVL树的验证
	bool IsAVLTree()
	{
		return _IsAVLTree(_pRoot);
	}

	// 右单旋
	void RotateR(Node* pParent)
	{
		Node* subL = pParent->_pLeft;
		Node* subLR = subL->_pRight;

		// 修改孩子指针指向，父亲指针指向
		pParent->_pLeft = subLR;
		if (subLR)
			subLR->_pParent = pParent;

		Node* parentParent = pParent->_pParent;

		subL->_pRight = pParent;
		pParent->_pParent = subL;


		// parent有可能是整棵树的根，也可能是局部的子树
	    // 如果是整棵树的根，要修改_root
	    // 如果是局部的指针要跟上一层链接
		if (parentParent == nullptr)
		{
			// 修改_root
			_pRoot = subL;
		}
		else
		{
			// 跟上一层链接
			if (pParent = parentParent->_pLeft)
			{
				parentParent->_pLeft = subL;
			}
			else
			{
				parentParent->_pRight = subL;
			}
			subL->_pParent = parentParent;
		}
		pParent->_bf = 0;
		subL->_bf = 0;
	}
	// 左单旋
	void RotateL(Node* pParent)
	{
		Node* subR = pParent->_pRight;
		Node* subRL = subR->_pLeft;

		// 修改孩子指针指向，父亲指针指向
		pParent->_pLeft = subRL;
		if (subRL)
			subRL->_pParent = pParent;

		Node* parentParent = pParent->_pParent;

		subR->_pLeft = pParent;
		pParent->_pParent = subR;


		// parent有可能是整棵树的根，也可能是局部的子树
		// 如果是整棵树的根，要修改_root
		// 如果是局部的指针要跟上一层链接
		if (parentParent == nullptr)
		{
			// 修改_root
			_pRoot = subR;
		}
		else
		{
			// 跟上一层链接
			if (pParent = parentParent->_pLeft)
			{
				parentParent->_pLeft = subR;
			}
			else
			{
				parentParent->_pRight = subR;
			}
			subR->_pParent = parentParent;
		}
		pParent->_bf = 0;
		subR->_bf = 0;
	}
	// 左右双旋
	void RotateLR(Node* pParent)
	{
		Node* subL = pParent->_pLeft;
		Node* subLR = subL->_pRight;
		int bf = subLR->_bf;

		RotateL(subL); // 左旋
		RotateR(pParent); // 右旋

		if (bf == 0)
		{
			subL->_bf = 0;
			subLR->_bf = 0;
			pParent->_bf = 0;
		}
		else if (bf == 1)
		{
			subL->_bf = -1;
			subLR->_bf = 0;
			pParent->_bf = 0;
		}
		else if (bf == -1)
		{
			subL->_bf = 0;
			subLR->_bf = 0;
			pParent->_bf = 1;
		}
		else
		{
			assert(false);
		}
	}
	// 右左双旋
	void RotateRL(Node* pParent)
	{
		Node* subR = pParent->_pRight;
		Node* subRL = pParent->_pLeft;
		int bf = subRL->_bf;

		RotateR(subR); // 右旋
		RotateL(pParent); // 左旋

		if (bf == 0)
		{
			subR->_bf = 0;
			subRL->_bf = 0;
			pParent->_bf = 0;
		}
		else if (bf == 1)
		{
			subR = 0;
			subRL->_bf = 0;
			pParent->_bf = -1;
		}
		else if(bf == -1)
		{
			subR->_bf = 1;
			subRL->_bf = 0;
			pParent->_bf = 0;
		}
		else
		{
			assert(false);
		}
	}

private:
	// 根据AVL树的概念验证pRoot是否为有效的AVL树
	int _Height(Node* root)
	{
		if (root == nullptr)
		{
			return 0;
		}

		int leftHeight = _Height(root->_pLeft);
		int rightHeight = _Height(root->_pRight);

		return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
	}
	bool _IsAVLTree(Node* pRoot)
	{
		// 空树也是AVL树
		if (pRoot == nullptr)
		{
			return true;
		}

		// 计算pRoot结点的平衡因子：即pRoot左右子树的高度差
		int leftHeight = _Height(pRoot->_pLeft);
		int rightHeight = _Height(pRoot->_pRight);
		int diff = rightHeight - leftHeight;

		if (diff >= 2 || diff <= -2)
		{
			cout << pRoot->_data << "高度差异常" << endl;
			return false;
		}
		if (pRoot->_bf != diff)
		{
			cout << pRoot->_data << "平衡因子异常" << endl;
			return false;
		}
		// pRoot的左和右如果都是AVL树，则该树一定是AVL树
		return _IsAVLTree(pRoot->_pLeft) && _IsAVLTree(pRoot->_pRight);
	}

private:
	Node* _pRoot;
};